/*  iImager (network disk imaging via iscsi)
	Copyright (C) 2009-2010 Yao Wei <njustyw@gmail.com>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include <stdlib.h>
#include <stdio.h>
#include "lib_imaging.h"
#include "trimps_forensic_global.h"
#include "cvTimer.h"
#include "Q_RWThread.h"
#include "Q_HDObject.h"


#define PROFILE_TIMECOST
#undef PROFILE_TIMECOST	

#define READ_BUFFERS 16
//#define READ_BUFFER_SIZE (8*1024*1024)
#define READ_BUFFER_SIZE (1024*1024) /*for machine having small memory*/
enum
{
	BUFFER_ON_FREE = 0x00,
	BUFFER_ON_VERIFY = 0x01,
	BUFFER_ON_VERIFY2 = 0x02,
	BUFFER_ON_WRITE = 0x04,
	BUFFER_ON_ERROR = 0x08,
	BUFFER_ON_DROP = 0x10,
	BUFFER_ON_DONE = 0x11
};

struct Q_MT_BUFFER* allocate_qbuffer_for_thread_copy()
{
	struct Q_MT_BUFFER* buffer = new struct Q_MT_BUFFER[READ_BUFFERS];
	if(!buffer) 	return NULL;

	for(int i = 0; i < READ_BUFFERS; i++)
	{
		buffer[i].aligned_buffer = new unsigned char[READ_BUFFER_SIZE];
		if(buffer[i].aligned_buffer == NULL)
		{
			for(int j = 0; j < i; j++)
				delete []buffer[j].aligned_buffer;
			delete []buffer;
			return NULL;
		}
	}

	return buffer;
}

void initialize_qbuffer_state(struct Q_MT_BUFFER* buffer, quint8 hashType)
{
	for(int i = 0; i < READ_BUFFERS; i++)
	{
		buffer[i].state =  BUFFER_ON_FREE;
		buffer[i].read_size = 0;
		buffer[i].hash.hash_type_mask = hashType;
	}
}

void release_qbuffer_for_thread_copy(struct Q_MT_BUFFER* buffer)
{
	if(buffer)
	{
		for(int i = 0; i < READ_BUFFERS; i++)
		{
			delete []buffer[i].aligned_buffer;
			buffer[i].aligned_buffer = 0;
		}

		delete []buffer;
		buffer = 0;
	}
}

struct Q_MT_LOCKER* allocate_qlocker_for_thread_copy()
{
	struct Q_MT_LOCKER* locker = new struct Q_MT_LOCKER[READ_BUFFERS];
	return locker;
}

void release_qlocker_for_thread_copy(Q_MT_LOCKER* locker)
{
	if(locker)
	{
		delete []locker;
		locker = 0;
	}
}

struct Q_MT_PARAM* allocate_param_for_thread_copy()
{
	struct Q_MT_PARAM* param = new struct Q_MT_PARAM;
	if(!param)
		return NULL;

	param->device = new struct device_handle;
	if(!param->device)
	{
		delete param;
		param = NULL;
		return NULL;
	}

	param->handle = new struct lib_imaging;
	if(!param->handle)
	{
		delete param->device;
		delete param;
		param = NULL;
		return NULL;
	}

	param->progress = new struct progress_t;
	if(!param->progress)
	{
		delete param->handle;
		delete param->device;
		delete param;
		delete param;
		return NULL;
	}
	memset(param->progress, 0, sizeof(*param->progress));

	param->isdrop.fetchAndStoreOrdered(0);
	param->ispause.fetchAndStoreOrdered(0);

	return param;
}

void release_param_for_thread_copy(struct Q_MT_PARAM* param)
{
	if(param)
	{
		delete param->device;
		delete param->handle;
		delete param->progress;
		delete param;
		param = 0;
	}
}

bool initialize_device_and_imaging_handle(Q_MT_PARAM* param, const QString& deviceName, const QString& fileName, 
										  int imagingType, int hashType)
{
	char error_msg[1024];
	uint8_t media_type = 0;
	quint64 total_bytes;
	unsigned int bytes_per_sector;
	int retry ;
	int max_try = 5;
	
	if(device_handle_initialize(param->device) < 0)
		return false; 
	if(device_handle_open_input(param->device, deviceName.toUtf8().data(), error_msg) < 0)
		return false;

	retry = 0;
	do{
		retry ++;
		if(retry == max_try) 
			return false;
	}while(device_handle_get_media_size(param->device, &total_bytes) < 0);

	retry = 0;
	do{
		retry ++;
		if(retry == max_try) 
			return false;
	}while(device_handle_determine_media_information(param->device) < 0);

	retry = 0;
	do{
		retry ++;
		if(retry == max_try) 
			return false;
	}while(device_handle_get_media_type(param->device, &media_type) < 0);

	retry = 0;
	do{
		retry ++;
		if(retry == max_try) 
			return false;
	}while(device_handle_get_bytes_per_sector(param->device, &bytes_per_sector) < 0);

	int sIdx = fileName.lastIndexOf("/");
	if(sIdx != -1)
	{
		QString path = fileName.left(sIdx);
		QDir dir;
		if(!dir.exists(path))	
			dir.mkdir(path);
	}

	//! step 1
	if(lib_imaging_initialize(param->handle, imagingType, hashType, error_msg) < 0)
		return false;

	//! step 2
	if(lib_imaging_open_output(param->handle, fileName.toLocal8Bit().data(), error_msg) < 0)
		return false;

	//! step 2.1
	if(imagingType == IMAGING_EWF)
	{
		if(lib_imaging_setting_ewf(param->handle, param->device, NO_COMPRESSION,"", "", "", "", "", error_msg) < 0)
			return false;
	}

	//! step 3
	if(device_handle_seek_offset(param->device, param->handle->imaging->write_raw_size, error_msg) < 0)
		return false;
	if(lib_imaging_seek_offset(param->handle, param->handle->imaging->write_size, error_msg) < 0)
		return false;

	return true;
}

bool cleanup_device_and_imaging_handle(Q_MT_PARAM* param)
{
	char error_msg[1024];
	
	device_handle_close(param->device, error_msg);
	
	lib_imaging_remove(param->handle, error_msg, REMOVE_ALL);
	lib_imaging_close(param->handle, error_msg);
	lib_imaging_free(param->handle, param->err_w_msg);

	return true;
}

QReadThread::QReadThread(QObject* parent /* = 0 */):QThread(parent)
{

}

QReadThread::~QReadThread()
{
}

void QReadThread::setParam(struct Q_MT_PARAM* parm)
{
	qparam = parm;
}

void QReadThread::run()
{
	struct Q_MT_PARAM* param = (struct Q_MT_PARAM*)qparam;
	struct Q_MT_BUFFER* buffer = param->buffer;
	struct lib_imaging* handle = param->handle;
	struct device_handle* device = param->device;
	struct progress_t* progress = param->progress;
	struct Q_MT_LOCKER* locker = param->locker;
	QMutex& data_mutex = param->cond_mutex; 
	QMutex& cond_mutex = param->cond_mutex; 
	QWaitCondition& pause_cond = param->pause_cond;
	HDObject* obj = param->obj;
	int percentage;
	int remaining_time, elapsed_time;
	double current_speed;
	int chunk_size = READ_BUFFER_SIZE;
	int is_done = 0;
	int i = 0;
	qint64 read_all = 0;

#ifdef PROFILE_TIMECOST	
//	QTime t1;
//	int totalTime = 0;
	MTimer t2;
	double totalTime2 = 0;
#endif

	DEBUG3_MSG("Reader started..." + obj->devname + param->device->device_name);

	if(handle->t == IMAGING_EWF)
	{
		int bytes_per_chunk	= handle->parameter->bytes_per_chunk;
		chunk_size = (READ_BUFFER_SIZE / bytes_per_chunk) * bytes_per_chunk;
	}
    
	status_start(progress, device->media_size, handle->imaging->elapsed_time, handle->imaging->write_raw_size);
	while(!is_done)
	{
		QMutexLocker rLocker(&locker[i].buffer_mutex);

		while (!(buffer[i].state == BUFFER_ON_FREE || buffer[i].state == BUFFER_ON_DROP || buffer[i].state == BUFFER_ON_ERROR) )
		{
			locker[i].read_cv.wait(rLocker.mutex());
		}

		cond_mutex.lock();
		if(param->ispause == 1)
		{
			QString loc = QString("(%1, %2)").arg(obj->rackIndex).arg(obj->locationIndex);
			DEBUG2_MSG(loc + "Read-Thread paused");

			status_pause(progress);
			obj->setState(HD_Paused);
			pause_cond.wait(&cond_mutex);
			obj->setState(HD_Imaging);

			DEBUG2_MSG(loc + "Read-Thread continued");
		}
		cond_mutex.unlock();

#ifdef PROFILE_TIMECOST	
		//t1.restart();
		//t1 = QTime::currentTime();
		t2.start();
#endif
		if(buffer[i].state == BUFFER_ON_DROP || param->isdrop == 1)
		{
			param->isdrop.fetchAndStoreOrdered(1);
			buffer[i].state = BUFFER_ON_DROP;
			locker[i].verify_cv.wakeAll();
			break;
		}
		else if(buffer[i].state == BUFFER_ON_ERROR)
		{
			locker[i].verify_cv.wakeAll();
			break;
		}

		buffer[i].read_size = device_handle_read_buffer(device, buffer[i].aligned_buffer, chunk_size, param->err_r_msg);

		if (buffer[i].read_size == -1 )
		{
			buffer[i].state = BUFFER_ON_ERROR;
			locker[i].verify_cv.wakeAll();
			DEBUG3_MSG("READ ERROR...");
			emit signal_exceptioned(obj, param->err_r_msg, param);
			break;;
		}
		else if (buffer[i].read_size == 0 )
		{
			buffer[i].state = BUFFER_ON_DONE;
			is_done = 1;
			device_handle_close(device, param->err_r_msg);
		}
		else
		{
			buffer[i].state = BUFFER_ON_VERIFY;
		}

		if(read_all >= 34000000000ll)
		{
			buffer[i].state = BUFFER_ON_DONE;
			is_done = 1;
		}

		read_all += buffer[i].read_size;

		locker[i].verify_cv.wakeAll();

#ifdef PROFILE_TIMECOST	
		//totalTime += t1.elapsed();
		totalTime2 += t2.getElapsed();
		//qDebug() << "R" << /*t1.msecsTo(QTime::currentTime()) <<*/ t2.getElapsed() << /*totalTime  <<*/ totalTime2;
#endif

		if(status_update(progress, read_all, &percentage, &current_speed, &remaining_time, &handle->imaging->elapsed_time))
		{
			elapsed_time = handle->imaging->elapsed_time;
			obj->mean_speed = current_speed;
			obj->elapsed_time = elapsed_time;
			obj->processed_bytes = read_all + progress->start_size;

			emit signal_updateStatus(obj, percentage, current_speed, remaining_time, elapsed_time);
		}

		i++;
		i = i % READ_BUFFERS;
	}

	QThread::yieldCurrentThread();

#ifdef PROFILE_TIMECOST	
	FILE* f  = fopen("profile-r.dat", "a+");
	fprintf(f, "Reader thread (%d, %d) time cost %g millsecs\n\n", obj->rackIndex, obj->locationIndex, totalTime2);
	fclose(f);
#endif
}

QVerifyThread::QVerifyThread(QObject* parent /* = 0 */):QThread(parent)
{

}

QVerifyThread::~QVerifyThread()
{
}

void QVerifyThread::setParam(struct Q_MT_PARAM* parm, quint8 hashMask)
{
	qparam = parm;
	hashTypeMask = hashMask;
}

void QVerifyThread::run()
{
	struct Q_MT_PARAM* param = (struct Q_MT_PARAM*)qparam;
	struct Q_MT_BUFFER* buffer = param->buffer;
	struct lib_imaging* handle = param->handle;
	HDObject* obj = param->obj;
	struct progress_t* progress = param->progress;
	struct Q_MT_LOCKER* locker = param->locker;
	QMutex& data_mutex = param->cond_mutex; 
	QMutex& cond_mutex = param->cond_mutex; 
	QWaitCondition& pause_cond = param->pause_cond;
	struct hash_content hash;
	int is_done = 0;
	int i = 0;
	uint8_t* buf;
	int verify_count;

	data_mutex.lock();
	lib_imaging_hash_init_all(handle, &handle->hash_cont);
	lib_imaging_hash_copy(hashTypeMask, &handle->hash_cont, &hash);
	data_mutex.unlock();

#ifdef PROFILE_TIMECOST	
	MTimer t2, t3;
	double totalTime2 = 0, totalTime3 = 0;
	t3.start();
#endif	

	DEBUG3_MSG("Verifier started...");

	while(!is_done)
	{
		QMutexLocker vLocker(&locker[i].buffer_mutex);

		while (!(buffer[i].state == BUFFER_ON_VERIFY || buffer[i].state == BUFFER_ON_DONE || buffer[i].state == BUFFER_ON_DROP || buffer[i].state == BUFFER_ON_ERROR) )
		{
			locker[i].verify_cv.wait(vLocker.mutex());
		}
#ifdef PROFILE_TIMECOST	
		//estimate the time that between threads for synchronize
		totalTime3 += t3.getElapsed();
#endif


		cond_mutex.lock();
		if(param->ispause == 1)
		{
			QString loc = QString("(%1, %2)").arg(obj->rackIndex).arg(obj->locationIndex);
			DEBUG2_MSG(loc + "Verify-Thread paused");

			status_pause(progress);
			obj->setState(HD_Paused);
			pause_cond.wait(&cond_mutex);
			obj->setState(HD_Imaging);

			DEBUG2_MSG(loc + "Verify-Thread continued");
		}
		cond_mutex.unlock();

#ifdef PROFILE_TIMECOST	
		t2.start();
#endif

		if (buffer[i].state == BUFFER_ON_DROP)
		{
			locker[i].write_cv.wakeAll();
			break;
		}
		else if (buffer[i].state == BUFFER_ON_ERROR)
		{
			locker[i].write_cv.wakeAll();
			break;
		}
		else if (buffer[i].state == BUFFER_ON_DONE)
		{
			lib_imaging_hash_final(hashTypeMask, handle, &hash);
			is_done = 1;
			
			data_mutex.lock();
			lib_imaging_hash_copy(hashTypeMask, &hash, &handle->hash_cont);
			data_mutex.unlock();

			locker[i].write_cv.wakeAll();
			
			break;
		}

		verify_count = buffer[i].read_size;
		buf = (uint8_t*)buffer[i].aligned_buffer;

		lib_imaging_hash_update(hashTypeMask, handle, &hash, buf, verify_count);

		data_mutex.lock();
		lib_imaging_hash_copy(hashTypeMask, &hash, &buffer[i].hash);
		data_mutex.unlock();

		
		buffer[i].state = BUFFER_ON_WRITE;
		locker[i].write_cv.wakeAll();
		

		i++;
		i = i % READ_BUFFERS;

#ifdef PROFILE_TIMECOST	
		totalTime2 += t2.getElapsed();
		t3.start();
#endif
	}

	QThread::yieldCurrentThread();

#ifdef PROFILE_TIMECOST	
	FILE* f  = fopen("profile-v.dat", "a+");
	fprintf(f, "Synchronize between threads (%d, %d) time cost %g millsecs\n", obj->rackIndex, obj->locationIndex, totalTime3);
	fprintf(f, "Verifier thread (%d, %d) time cost %g millsecs\n\n", obj->rackIndex, obj->locationIndex, totalTime2);
	fclose(f);
#endif
}

QWriteThread::QWriteThread(QObject* parent /* = 0 */):QThread(parent)
{

}

QWriteThread::~QWriteThread()
{
}

void QWriteThread::setParam(struct Q_MT_PARAM* parm)
{
	qparam = parm;
}

void QWriteThread::run()
{
	struct Q_MT_PARAM* param = (struct Q_MT_PARAM*)qparam;
	struct Q_MT_BUFFER* buffer = param->buffer;
	struct lib_imaging* handle = param->handle;
	struct progress_t* progress = param->progress;
	struct Q_MT_LOCKER* locker = param->locker;
	QMutex& data_mutex = param->cond_mutex; 
	QMutex& cond_mutex = param->cond_mutex; 
	QWaitCondition& pause_cond = param->pause_cond;
	HDObject* obj = param->obj;
	int total_time;
	double mean_speed;
	int is_done = 0;
	int i = 0;
	int write_count_once, write_count_all, write_count;
	uint8_t* buf;

#ifdef PROFILE_TIMECOST	
	MTimer t2, t3;
	double totalTime2 = 0, totalTime3 = 0;
	t3.start();
#endif	

	DEBUG3_MSG("Writier started...");

	while(!is_done)
	{
        QMutexLocker wLocker(&locker[i].buffer_mutex);

		while (!(buffer[i].state == BUFFER_ON_WRITE || buffer[i].state == BUFFER_ON_DONE|| buffer[i].state == BUFFER_ON_DROP|| buffer[i].state == BUFFER_ON_ERROR) )
		{
			locker[i].write_cv.wait(wLocker.mutex());
		}

#ifdef PROFILE_TIMECOST	
		//estimate the time that between threads for synchronize
		totalTime3 += t3.getElapsed();
#endif

		cond_mutex.lock();
		if(param->ispause == 1)
		{
			QString loc = QString("(%1, %2)").arg(obj->rackIndex).arg(obj->locationIndex);
			DEBUG2_MSG(loc + "Write-Thread paused");

			status_pause(progress);
			obj->setState(HD_Paused);
			pause_cond.wait(&data_mutex);
			obj->setState(HD_Paused);

			DEBUG2_MSG(loc + "Write-Thread continued");
		}
		cond_mutex.unlock();

#ifdef PROFILE_TIMECOST	
		t2.start();
#endif

		if (buffer[i].state == BUFFER_ON_DROP)
		{
			locker[i].read_cv.wakeAll();
			emit signal_cancel(obj, param);
			break;
		}
		else if (buffer[i].state == BUFFER_ON_ERROR)
		{
			locker[i].read_cv.wakeAll();
			break;
		}
		else if (buffer[i].state == BUFFER_ON_DONE )
		{
			if(lib_imaging_finalize(handle, param->err_w_msg) < 0)
			{
				locker[i].read_cv.wakeAll();
				buffer[i].state = BUFFER_ON_ERROR;
				emit signal_exceptioned(obj, param->err_w_msg, param);
				break;
			}
			//! step 5
			lib_imaging_close(handle, param->err_w_msg);

			status_stop(progress, &mean_speed, &total_time);			
		
			emit signal_updateStatus(obj, 100, mean_speed, 0, handle->imaging->elapsed_time);
			
			is_done = 1;

			break;
		}

		write_count_all = 0;
		buf = (uint8_t*)buffer[i].aligned_buffer;
		write_count_once = DEFAULT_BYTES_PER_CHUNK;
		if(handle->t == IMAGING_EWF)
			write_count_once = handle->parameter->bytes_per_chunk;

		bool error = false;
		do
		{
			if(write_count_once > (buffer[i].read_size-write_count_all))
				write_count_once = buffer[i].read_size-write_count_all;

			write_count = lib_imaging_write_chunk(handle, buf, write_count_once, param->err_w_msg);

			if(write_count < 0)
			{
				locker[i].read_cv.wakeAll();
				emit signal_exceptioned(obj, param->err_w_msg, param);
				error = true;
				break;
			}

			write_count_all += write_count_once;
			buf += write_count_once;
		}while(write_count_all != buffer[i].read_size);
		if(error)
			break;

		data_mutex.lock();
		lib_imaging_backup_update(handle, &buffer[i].hash);
		data_mutex.unlock();

		buffer[i].state = BUFFER_ON_FREE;
		locker[i].read_cv.wakeAll();

		i++;
		i = i % READ_BUFFERS;

#ifdef PROFILE_TIMECOST	
		totalTime2 += t2.getElapsed();  
		t3.start();
#endif
	}

#ifdef PROFILE_TIMECOST	
	FILE* f  = fopen("profile-w.dat", "a+");
	fprintf(f, "Synchronize between threads (%d, %d) time cost %g millsecs\n", obj->rackIndex, obj->locationIndex, totalTime3);
	fprintf(f, "Writier thread (%d, %d) time cost %g millisec\n\n", obj->rackIndex, obj->locationIndex, totalTime2);
	fclose(f);
#endif

	QThread::yieldCurrentThread();

	if(is_done)
	{
		emit signal_finished(obj, mean_speed, total_time, param);
	}
}
